texture ImageTexture;
float2 Delta;
int Adjustment;
int SpatialRadius;
float ColorDistance;
float MinimumShifted;
int MaximumIteration;

sampler2D Image = sampler_state
{
  AddressU = Clamp;
  AddressV = Clamp;
  Texture = (ImageTexture);
};

struct Input
{
    float4 Position : POSITION0;
    float2 TexCoord0 : TEXCOORD0;
};

Input VS_Main(Input input)
{
	Input output;	 
	output.Position = float4(sign(input.Position.xy), 0.0f, 1.0f);	 
	output.TexCoord0 = input.TexCoord0 + Adjustment * Delta * 0.125;
	return output;
}

float3 GetYIQFromRGB(float3 rgb)
{
    rgb = rgb * 256.0f;

    float3 yiq = float3(0, 0, 0);
    // yiq.x = 0.2990f * rgb.r + 0.5870f * rgb.g + 0.1140f * rgb.b; // Y
    // yiq.y = 0.5957f * rgb.r - 0.2744f * rgb.g - 0.3212f * rgb.b; // I
    // yiq.z = 0.2114f * rgb.r - 0.5226f * rgb.g + 0.3111f * rgb.b; // Q
    
    return rgb;
}

float3 GetRGBFromYIQ(float3 yiq)
{
    float3 rgb = float3(0, 0, 0);
    // rgb.r = yiq.x + 0.9563f * yiq.y + 0.6210f * yiq.z; // R
    // rgb.g = yiq.x - 0.2721f * yiq.y - 0.6473f * yiq.z; // G
    // rgb.b = yiq.x - 1.1070f * yiq.y + 1.7046f * yiq.z; // B
    
    rgb = yiq / 256.0f;
    
    return rgb;
}

float GetSquaredEuclideanDistance(float3 a, float3 b)
{
    float x = a.x - b.x;
    float y = a.y - b.y;
    float z = a.z - b.z;

    float euclideanDistance = x * x + y * y + z * z;
    
    return euclideanDistance;
}

float GetSquaredEuclideanDistance(float2 a, float2 b)
{
    float x = a.x - b.x;
    float y = a.y - b.y;

    float euclideanDistance = x * x + y * y;
    
    return euclideanDistance;
}

float4 PS_MeanShift(Input input) : COLOR0 
{
    float SquaredSpatialRadius = SpatialRadius * SpatialRadius;
    float SquaredColorDistance = ColorDistance * ColorDistance;
    
    float4 currentColor = tex2D(Image, input.TexCoord0);
    
    float2 position = input.TexCoord0;
    
    float3 yiq = GetYIQFromRGB(currentColor.rgb);

    float shifted = 0.0f;    
    
    int iteration = 0;
        
    do
    {
        float2 lastPosition = position;
        float3 lastYIQ = yiq;

        float2 summedPosition = float2(0.0f, 0.0f);
        float3 summedYIQ = float3(0.0f, 0.0f, 0.0f);
        
        int count = 0;
      
        for (int v = -SpatialRadius; v <= SpatialRadius; v++)
        {            
            for (int u = -SpatialRadius; u <= SpatialRadius; u++)
            {                
                if (v * v + u * u <= SquaredSpatialRadius)
                {
                    float x = max(0.0f, min(1.0f, input.TexCoord0.x + u * Delta.x));
                    float y = max(0.0f, min(1.0f, input.TexCoord0.y + v * Delta.y));
                    
                    currentColor = tex2Dlod(Image, float4(x, y, 0.0f, 0.0f));
                                        
                    float3 currentYIQ = GetYIQFromRGB(currentColor.rgb);
                    
                    float e = GetSquaredEuclideanDistance(yiq, currentYIQ);

                    if (e <= SquaredColorDistance)
                    {
                        summedPosition.x += u;
                        summedPosition.y += v;
                        
                        summedYIQ += currentYIQ;
                        
                        count++;
                    }
                }
            }
        }

        float invertCount = 1.0f / count;
        
        yiq = summedYIQ * invertCount;
        
        position = summedPosition * invertCount + 0.5f;
        
        shifted =
            GetSquaredEuclideanDistance(position, lastPosition) +
            GetSquaredEuclideanDistance(yiq, lastYIQ);
            
        iteration++;
    }
    while (shifted > MinimumShifted && iteration < MaximumIteration);
    
    float3 rgb = GetRGBFromYIQ(yiq);

  	return float4(rgb, currentColor.w); 
}

/*
float4 PS_OptimizedMeanShift(Input input) : COLOR0 
{
    float SquaredSpatialRadius = SpatialRadius * SpatialRadius;
    float SquaredColorDistance = ColorDistance * ColorDistance;
    
    float4 currentColor = tex2D(Image, input.TexCoord0);
    
    float2 position = input.TexCoord0;
    
    float3 yiq = GetYIQFromRGB(currentColor.rgb);

    float shifted = 0.0f;    
    
    int iteration = 0;
        
    do
    {
        float2 lastPosition = position;
        float3 lastYIQ = yiq;

        float2 summedPosition = float2(0.0f, 0.0f);
        float3 summedYIQ = float3(0.0f, 0.0f, 0.0f);
        
        int count = 0;
      
        for (int v = -SpatialRadius; v <= SpatialRadius; v++)
        {            
            for (int u = -SpatialRadius; u <= SpatialRadius; u++)
            {                
                if (v * v + u * u <= SquaredSpatialRadius)
                {
                    float x = max(0.0f, min(1.0f, position.x + u * Delta.x));
                    float y = max(0.0f, min(1.0f, position.y + v * Delta.y));
                    
                    currentColor = tex2Dlod(Image, float4(x, y, 0.0f, 0.0f));
                                        
                    float3 currentYIQ = GetYIQFromRGB(currentColor.rgb);
                    
                    float e = GetSquaredEuclideanDistance(yiq, currentYIQ);

                    if (e <= SquaredColorDistance)
                    {
                        summedPosition.x += u;
                        summedPosition.y += v;
                        
                        summedYIQ += currentYIQ;
                        
                        count++;
                    }
                }
            }
        }

        float invertCount = 1.0f / count;
        
        yiq = summedYIQ * invertCount;
        
        position = summedPosition * invertCount + 0.5f;
        
        shifted =
            GetSquaredEuclideanDistance(position, lastPosition) +
            GetSquaredEuclideanDistance(yiq, lastYIQ);
            
        iteration++;
    }
    while (shifted > MinimumShifted && iteration < MaximumIteration);
    
    float3 rgb = GetRGBFromYIQ(yiq);

  	return float4(rgb, currentColor.w); 
}
*/

technique MeanShift
{
    pass Pass0
    {        
		VertexShader = compile vs_3_0 VS_Main();
        PixelShader = compile ps_3_0 PS_MeanShift();
    }
}

/*
technique OptimizedMeanShift
{
    pass Pass0
    {        
		VertexShader = compile vs_3_0 VS_Main();
        PixelShader = compile ps_3_0 PS_OptimizedMeanShift();
    }
}
*/